Age of Rust

Python versions on Alliance clusters

Use module

module spider python

module spider python/3.12.4

module load python/3.12.4

Python packages on Alliance clusters

Create a Python virtual environment:

python -m venv ~/env

Activate it:

source ~/env/bin/activate

Update pip from wheel:

python -m pip install --upgrade pip --no-index

Use pip to install packages and --no-index to use wheels whenever possible:

python -m pip install --no-index jax[cuda12] jax-ai-stack[grain]

Install uv

Help

List of commands and options:

uv

List of options:

uv <command> -h    # e.g. uv init -h

Man page:

uv help <command>  # e.g. uv help init

Drop-in replacement

You can add uv in front of your usual venv and pip commands

This actually runs uv (and neither pip nor venv) so you get the speedup, but it keeps everything compatible

Create a virtual env

uv venv

uv venv --python 3.12

By default, the virtual env is called .venv. If you don’t change its name, uv will use it automatically so you don’t need to source it

Install packages in virtual env

uv pip install jax flax

uv pip install "git+https://github.com/jax-ml/jax"
uv pip install "git+https://github.com/jax-ml/jax@main"
uv pip install "git+https://github.com/jax-ml/jax@766e68c4813a30e29b4fcefaa3253a42d0e197be"

uv pip install -r requirements.txt
uv pip install -r pyproject.toml

All your usual commands work

uv pip uninstall jax
uv pip list
uv pip freeze
…

Install Python

uv python install

uv python install 3.12.3
uv python install '>=3.8,<3.10'

uv python install cpython

Manage versions

uv python list

uv python uninstall 3.13

Initialize projects

uv init my_project

uv init --python 3.12 my_project

uv init --no-readme --no-description

Project structure

eza -aT my_project

my_project
├── .python-version
├── main.py
├── pyproject.toml
└── README.md

bat -p my_project/pyproject.toml

[project]
name = "demo"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
    "matplotlib>=3.10.1",
    "polars>=1.29.0",
]

Add dependencies

You need to cd into the project, then you can add dependencies:

cd my_project
uv add polars matplotlib

This creates a virtual env called .venv and a uv.lock:

eza -aTL 1

my_project
├── .python-version
├── .venv
├── main.py
├── pyproject.toml
├── README.md
└── uv.lock

Here again, no need to source the virtual env as long as you use uv

Project file

Gets populated automatically with dependencies:

bat -p pyproject.toml

[project]
name = "demo"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
    "matplotlib>=3.10.1",
    "polars>=1.29.0",
]

List explicitly installed dependencies

uv tree -d 1

Resolved 13 packages in 0.74ms
demo v0.1.0
├── matplotlib v3.10.1
└── polars v1.29.0

List all dependencies

uv pip list

Package         Version
--------------- -----------
contourpy       1.3.2
cycler          0.12.1
fonttools       4.57.0
kiwisolver      1.4.8
matplotlib      3.10.1
numpy           2.2.5
packaging       25.0
pillow          11.2.1
polars          1.29.0
pyparsing       3.2.3
python-dateutil 2.9.0.post0
six             1.17.0

Manage dependencies

Update all dependencies in lock file and virtual env:

uv sync -U

Remove dependencies:

uv remove matplotlib